import copy
from typing import Dict, Iterator, List, Optional, Union
# 代理类
from qwen_agent import Agent
# 助理类
from qwen_agent.agents import Assistant
# 基础聊天模型
from qwen_agent.llm import BaseChatModel
from qwen_agent.llm.schema import ContentItem, Message
from qwen_agent.tools import BaseTool
from qwen_agent.gui import WebUI


# 继承自Agent基类
class VisualStorytelling(Agent):
    # function_list：可以使用的工具列表，默认为空。
    # llm：指定的语言模型，默认为空。
    def __init__(self,
                 function_list: Optional[List[Union[str, Dict,
                 BaseTool]]] = None,
                 llm: Optional[Union[Dict, BaseChatModel]] = None):
        super().__init__(llm=llm)

        # 初始化视觉理解Agent
        self.image_agent = Assistant(llm={'model': 'qwen-vl-max'})

        # 初始化写作Agent
        # 提供指定工具和文件，用于支持写作任务。
        self.writing_agent = Assistant(
            llm=self.llm,
            function_list=function_list,
            system_message='你扮演一个想象力丰富的学生，你需要先分析图片内容，并把分析的图片内容放在开头，再参考知识库中教你的写作技巧，发挥你的想象力，写一篇800字的记叙文。',
            files=['./resource/article_tutorial.md'])

    # 定义工作流程，用于处理一组输入消息，输出最终作文。
    def _run(self,
             messages: List[Message],
             lang: str = 'zh',
             max_ref_token: int = 4000,
             **kwargs) -> Iterator[List[Message]]:
        assert isinstance(messages[-1]['content'], list) and any([
            item.image for item in messages[-1]['content']
        ]), 'This agent requires input of images'

        # 图片理解
        new_messages = copy.deepcopy(messages)
        new_messages[-1]['content'].append(
            ContentItem(text='请详细描述这张图片的所有细节内容'))
        response = []
        for rsp in self.image_agent.run(new_messages):
            yield response + rsp
        response.extend(rsp)
        new_messages.extend(rsp)

        # 写文章
        new_messages.append(Message('user', '开始根据以上图片内容编写你的记叙文吧！'))
        for rsp in self.writing_agent.run(new_messages,
                                          lang=lang,
                                          max_ref_token=max_ref_token,
                                          **kwargs):
            yield response + rsp


# 配置您所使用的 LLM。
llm_cfg = {
    'model': 'qwen-max',
    'model_server': 'dashscope'
}

bot = VisualStorytelling(llm=llm_cfg)
image = './resource/article.png'
messages = [Message('user', [ContentItem(image=image)])]
messages[-1]['content'].append(ContentItem(text='看图写作文'))

# WebUI(bot).run()

# 追踪前一次的内容以避免重复
prev_content = ""
prev_response_len = 0
for response in bot.run(messages):
    # 检查响应数组长度是否增加
    current_response_len = len(response)
    if current_response_len > prev_response_len:
        # 新元素被添加，打印其完整内容
        latest_msg = response[-1]
        if latest_msg.content != '':
            print(latest_msg.content, end='')
            prev_content = latest_msg.content
    elif current_response_len == prev_response_len and response:
        # 检查最后一个元素的内容是否更新
        latest_msg = response[-1]
        if latest_msg.content != prev_content:

            # 只打印新的内容
            print(latest_msg.content[len(prev_content):], end='')
            prev_content = latest_msg.content
    
    prev_response_len = current_response_len
